import os
import sys
from cffi import FFI
from misc import generate_extension_name
import numpy as np
from timer import timer
ffibuilder = FFI()

# 1. 声明C函数接口
cdef = """
void broadcast_add_3d(
    float* a_data, int a_shape[3], int a_stride[3],
    float* b_data, int b_shape[3], int b_stride[3],
    float* c_data, int c_shape[3], int c_stride[3]);
"""
ffibuilder.cdef(cdef)

mod_name = "_tensor_ops"
platform = sys.platform
if 'win' in platform:
    compile_args = ['/O2', '/arch:AVX2']
else:
    compile_args = ['-O3', '-mavx2', '-march=native']
# 2. 指定C源码（支持内联代码或外部文件）
ffibuilder.set_source(mod_name,  # 生成的模块名
    cdef,
    libraries=[],  # 依赖的外部库
    sources=['broadcast_add.c'],  # 引用外部C文件
    extra_compile_args=compile_args,  # 添加编译优化参数
)
ext_name = generate_extension_name(mod_name)
# if not os.path.exists(ext_name):
ffibuilder.compile(verbose=True)  # 自动触发编译
    
from _tensor_ops import ffi, lib

        
def broadcast_add_wrapper(a, b, backend):
    """基于C实现的广播加法，兼容Numpy广播规则"""
    # 验证广播兼容性
    try:
        np.broadcast_shapes(a.shape, b.shape)
    except ValueError as e:
        raise ValueError(f"广播形状不兼容: {a.shape} vs {b.shape}") from e
    
    # 转换为C兼容的三维形状和步长
    def to_3d(arr):
        shape = list(arr.shape) + [1]*(3-len(arr.shape))
        strides = [s//arr.itemsize for s in arr.strides] + [0]*(3-len(arr.strides))
        return shape, strides
    a_shape, a_strides = to_3d(a)
    b_shape, b_strides = to_3d(b)
    out_shape = np.broadcast_shapes(a_shape, b_shape)
    # 创建输出数组（确保内存连续）
    c = np.empty(out_shape, dtype=np.float32, order='C')
    if backend == 'np':
        np_broadcast_add(a, b, c)
    else:
        out_strides = [s//4 for s in c.strides]
        a = np.asarray(a, dtype=np.float32, order='C')
        b = np.asarray(b, dtype=np.float32, order='C')
        c_broadcast_add(a, b, c, a_shape, b_shape, out_shape, a_strides, b_strides, out_strides)
    return c

@timer(repeats=10, warm_up=2)
def np_broadcast_add(a, b, c):
    np.add(a, b, out=c)
    
@timer(repeats=10, warm_up=2)
def c_broadcast_add(a, b, c, a_shape, b_shape, out_shape, a_strides, b_strides, out_strides):
    # 调用C函数
    a_ptr = ffi.cast("float*", a.ctypes.data)
    b_ptr = ffi.cast("float*", b.ctypes.data)
    c_ptr = ffi.cast("float*", c.ctypes.data)
    lib.broadcast_add_3d(
        a_ptr,
        ffi.new('int[3]', a_shape),
        ffi.new('int[3]', a_strides),
        b_ptr,
        ffi.new('int[3]', b_shape),
        ffi.new('int[3]', b_strides),
        c_ptr,
        ffi.new('int[3]', out_shape),
        ffi.new('int[3]', out_strides)  # 转换为元素步长
    )